'''
绘制loss曲线
参考链接：https://blog.csdn.net/qq_32855463/article/details/86610098

'''

import numpy as np
import matplotlib.pyplot as plt
import pylab as pl
from tqdm import tqdm, trange
from mpl_toolkits.axes_grid1.inset_locator import inset_axes


total_loss = []
loss_cls = []
loss_box_reg = []
iter = []
isFirst=1
with open("YOLOF\\220311\\loss.txt", "r") as f:
    for line in tqdm(f.readlines()):
        line = line.strip('\n')  #去掉列表中每一个元素的换行符
        if("Saving checkpoint" in line):
            continue
        if (int(line.split("iter: ")[1].split("/320000")[0]) > 100000):
            break
        iter.append(int(line.split("iter: ")[1].split("/320000")[0]))

        total_loss.append(float(line.split("total_loss: ")[1].split("  loss_cls: ")[0]))
        loss_cls.append(float(line.split("  loss_cls: ")[1].split("  loss_box_reg")[0]))
        if(isFirst==1):
            loss_box_reg.append(float(line.split("loss_box_reg: ")[1].split("      data_time:")[0]))
            isFirst=2
        else:
            loss_box_reg.append(float(line.split("loss_box_reg: ")[1].split("    time:")[0]))
        #print(line)



fig = plt.figure(figsize = (20,20))       #figsize是图片的大小`



p1=pl.plot(iter,total_loss,'g-',label=u'total_loss')
pl.legend()
# ‘’g‘’代表“green”,表示画出的曲线是绿色，“-”代表画的曲线是实线，可自行选择，label代表的是图例的名称，一般要在名称前面加一个u，如果名称是中文，会显示不出来，目前还不知道怎么解决。
p2 = pl.plot(iter, loss_cls,'r-', label = u'loss_cls')
pl.legend()
#显示图例
p3 = pl.plot(iter,loss_box_reg, 'b-', label = u'loss_box_reg')
pl.legend()
pl.xlabel(u'iters')
pl.ylabel(u'loss')
plt.title('Compare loss for different models in training')
#plt.savefig("YOLOF\\220311\\lossFor100000.png")
plt.show()

